# 高斯光束电场函数
import numpy as np
def electric(wavelength, omega, phase, fwhm, tau, E_peak, waist, rho, r, z, t):
	sigma = fwhm /(2 * np.sqrt(2*np.log(2)))
	k0 = 2 * np.pi / wavelength
	zr = np.pi * waist**2 / wavelength
	if z==rho:
		w_z = waist
		r_z = 1
		eta_z = 0
	else:
		w_z = waist * np.sqrt(1. + np.power((z-rho)/zr,2))
		r_z = (z - rho)*(1 + np.power(zr/(z-rho), 2))
		eta_z = np.arctan2(z-rho, zr)

	arg = 0
	arg += omega*t + phase
	arg += -k0 * (r**2) / (2 * r_z)
	arg += eta_z

	Er_envelop = (waist/w_z) * np.exp(-(r/w_z)**2)
	Et_envelop = np.exp(-0.5*((t-tau)/sigma)**2)

	EE = E_peak * Er_envelop * Et_envelop * np.cos(arg)
	return EE

def Et_envelop(t,tau,fwhm, E_peak):
	sigma = sigma = fwhm /(2 * np.sqrt(2*np.log(2)))
	return E_peak*np.exp(-0.5*((t-tau)/sigma)**2)